from collections import defaultdict
import warnings

from histogram.h1d import Histogram1D
from histogram.h2d import Histogram2D
from histogram.h3d import Histogram3D

class DuplicateKeyError(Exception):
    pass

class HistogramCollection(object):
    def __init__(self, hl=[], module_name=None):
        self._hd    = {}
        self._names = []
        self._assoc_cols = defaultdict(set)

        for h in hl:
            self.add(h)

        self.colfilters = {}
        self.module_name = module_name        

    def add(self, hist):
        if hist.name in self._hd:
            msg = "Histogram with name `{0}' is already present"
            raise DuplicateKeyError(msg.format(hist.name))
                
        self._hd[hist.name] = hist
        self._names.append(hist.name)
        self._names.sort()

        for col in hist.assoc_cols:
            self._assoc_cols[col].add(hist)

    def assoc_cols(self):
        return set(self._assoc_cols.keys())

    def col_used_by(self, col):
        return self._assoc_cols[col]        

    def fill(self, coldict):
        missing_cols = self.assoc_cols() - set(coldict.keys())
        if len(missing_cols) > 0:
            msg = "Column dict is missing columns {}; "
            msg += "not all histograms will be filled"
            warnings.warn(msg)

        for h in self.iter1d():
            ac = h.assoc_cols[0]
            if h.name in self.colfilters:
                h.fill(self.colfilters[h.name](coldict[ac]))
            else:
                h.fill(coldict[ac])

        for h in self.iter2d():
            ac1, ac2 = h.assoc_cols
            if h.name in self.colfilters:
                h.fill(self.colfilters[h.name](coldict[ac1], coldict[ac2]))
            else:
                h.fill((coldict[ac1], coldict[ac2]))
        
        for h in self.iter3d():
            ac1, ac2, ac3 = h.assoc_cols
            if h.name in self.colfilters:
                h.fill(self.colfilters[h.name](coldict[ac1], 
                                               coldict[ac2],
                                               coldict[ac3]))
            else:
                h.fill((coldict[ac1], coldict[ac2], coldict[ac3]))
                                       

    def column_fillable(self, tc):
        for n in self._names:
            _assoc_cols = set(self._hd[n]._assoc_cols)
            if _assoc_cols.issubset(tc):
                yield self._hd[n]    
                    
    def not_column_fillable(self, tc):
        for n in self._names:
            _assoc_cols = set(self._hd[n]._assoc_cols)
            if not _assoc_cols.issubset(tc):
                yield self._hd[n]

    def names(self):
        return self._names[:]
            
    def iter1d(self):
        for n in self._names:
            if type(self._hd[n]) == Histogram1D:
                yield self._hd[n]

    def iter2d(self):
        for n in self._names:
            if type(self._hd[n]) == Histogram2D:
                yield self._hd[n]

    def iter3d(self):
        for n in self._names:
            if type(self._hd[n]) == Histogram3D:
                yield self._hd[n]            

    def __setitem__(self, nm, hist):
        if nm in self:
            self._hd[nm] = hist
        elif nm == hist.name:
            self.add(hist)
        else:
            msg = "Key (`{}') does not match histogram name (`{}')"
            raise NotImplementedError(msg.format(nm, hist.name))
         
    def __getitem__(self, key):
        if type(key) == type('howdy!'):
            return self._hd[key]
        
        return self._hd[self._names[key]]

    def __iter__(self):
        for n in self._names:
            yield self._hd[n]            

    def __len__(self):
        return len(self._names)

    def __contains__(self, key):
        if type(key) == type("howdy!"):
            return key in self._names
        else:
            return key in self._hd.values()
